#!/usr/bin/env python

#
# Please see the LICENSE file included in this distribution.
#

"""Gmail Backup Utility.

This module contains a number of classes and stand alone functions to
implement an IMAP-based Gmail backup utility.  The classes are
IMAP4_SSL_STATS (which extends IMAP_SSL with basic transfer statistics),
gmailConnection (which provides an abstract interface for interacting
with a Gmail account), gmailBackup (which uses gmailConnection to
perform backup and restore operations), and logger (which managers the
display of status information as these classes are doing their work).
This module also contains stand-alone usage() and main() functions,
providing a complete command-line tool for Gmail backup.  See usage()
for information on using this tool.
"""

import datetime
import fnmatch
import getopt
import getpass
import hashlib
import imaplib
import locale
import os
import re
import socket
import string
import sys
import time

# The first part (before the hyphen) is the software version number and
# the second part is the backup storage format version number.  If either
# of the first two numbers in the backup storage format version number
# differ, the format is incompatible.
version = '00.01.02-00.02.00'

# todo
# - do we really need uids?
# - locking to prevent multiple concurrent backups/restores
# - ensure that can tolerate new arrivals during backup

#
# REGULAR EXRPESSION OBJECTS
# 

mailboxRE = re.compile(r'\((?P<flags>.*?)\) "(?P<delimiter>.*)" (?P<name>.*)')
noSelectRE = re.compile(r'.*\\Noselect.*')
noBkupMailboxRE = re.compile(r'^\[Gmail\]/(Spam)|(Trash)|(All Mail)$')
gmailMailboxRE = re.compile(r'^(\[Gmail\].*)|(INBOX)$')
uidRE = re.compile(r'.*UID (?P<uid>[0-9]+).*')
uidNextRE = re.compile(r'.*UIDNEXT (?P<uid>[0-9]+).*')
uidValidityRE = re.compile(r'.*UIDVALIDITY (?P<uid>[0-9]+).*')
indexFileRE = re.compile(r'^index_(?P<date>[^_]*)_(?P<uidn>[^_]*)_(?P<uidv>[^.]*)(?P<ext>\.bak)?$')
idateRE = re.compile(r'.*(?P<idate>INTERNALDATE \"[^\"]+\").*')
msgIdRE = re.compile(r'^Message-Id: *<?(?P<id>[^>]+)>?', re.IGNORECASE)
emlFileRE = re.compile(r'^(?P<base>(?P<date>..............)_(?P<msgid>.*))\.eml$')
appendRE = re.compile(r'^\[APPENDUID [^ ]+ (?P<id>[^ ]+)\] \(Success\)')


###
### CLASS DEFINITIONS
###

#
# class IMAP4_SSL_STATS
#

class IMAP4_SSL_STATS(imaplib.IMAP4_SSL):
    """IMAP4 client class over SSL connection with transfere statistics

    Instantiate with: IMAP4_SSL_STATS([host[, port[, keyfile[, certfile]]]])

      host - host's name (default: localhost)
      port - port number (default: standard IMAP4 SSL port)
      keyfile - PEM formatted file that contains private key (default: None)
      certfile - PEM formatted certificate chain file (default: None)

    IMAP4_SSL_STATS adds a getStats() method to IMAP4_SSL, returning
    the number of bytes send and received since the object was
    instantiated.

    See the docstring of the parent class for more information.
    """

    def __init__(self, host='', port=imaplib.IMAP4_SSL_PORT, keyfile=None,
                 certfile=None):
        self.bytesSent = 0
        self.bytesRecv = 0
        imaplib.IMAP4_SSL.__init__(self, host, port, keyfile, certfile)

    def getStats(self):
        """Return number of bytes sent and received since instantiation."""
        return self.bytesSent,self.bytesRecv

    def read(self, size):
        """Call parent implementation and record number of bytes received."""
        result = imaplib.IMAP4_SSL.read(self, size)
        self.bytesRecv += len(result)
        return result

    def readline(self):
        """Call parent implementation and record number of bytes received."""
        result = imaplib.IMAP4_SSL.readline(self)
        self.bytesRecv += len(result)
        return result
            
    def send(self, data):
        """Record number of bytes sent and call parent implementation."""
        self.bytesSent += len(data)
        result = imaplib.IMAP4_SSL.send(self, data)
        return result


class gmailConnection:
    """Class for managing a Gmail IMAP connection

    Instantiate with: gmailConnection(user, password, logger)

      user - Gmail account user name (e.g., user@gmail.com)
      password - Gmail account password
      logger - logger object (defined in this module)
      host - Gmail IMAP host (default: imap.googlemail.com)

    Although gmailConnection uses IMAP to connect to Gmail servers,
    it exposes Gmail abstractions (e.g., labels) instead of IMAP
    abstractions (e.g., mailboxes).  This class should only be used
    to connect to a Gmail IMAP server, for it makes Gmail-specific
    assumptions (e.g., that there exists a mailbox called
    '[Gmail]/All Mail').
    """

    defaultHost = 'imap.googlemail.com'
    host = defaultHost
    allMail = '[Gmail]/All Mail'
    trashMail = '[Gmail]/Trash'
    dtFormat = '%Y%m%d%H%M%S'
    currentMailbox = None

    #
    # Public methods
    #

    def __init__(self, user, password, logger, host=None):
        if host:
            self.host = host
        self.l = logger
        try:
            self.s = IMAP4_SSL_STATS(self.host)
        except socket.gaierror,e:
            # e is an (errno,description) pair
            raise Exception("Unable to connect to host '%s' (%s)" %
                            (self.host, e[1]))
        try:
            self.s.login(user, password)
        except imaplib.IMAP4_SSL.error,e:
            raise Exception("Unable to login to host (%s)" % e)
        self._selectAll()

    def getStats(self):
        """Return bytes sent and received since instantiation."""
        return self.s.getStats()

    def getUidNV(self):
        """Return UIDNEXT and UIDVALITIDY.  See module comments."""
        self._selectAll()
        typ,data = self.s.status(self.allMail, '(UIDNEXT UIDVALIDITY)')
        self._check('STATUS', typ, data)
        results = string.replace(data[0], '"%s" ' % (self.allMail), '')
        uidNext = int(uidNextRE.match(results).group(1))
        uidValidity = int(uidValidityRE.match(results).group(1))
        return (uidNext, uidValidity)

    def getLabels(self):
        """Return list of labels in account."""
        self._selectAll()
        typ,data = self.s.list('')
        self._check('LIST', typ, data)
        labs = map(lambda x: mailboxRE.match(x).group(3).strip('"'),
                   filter(lambda x: not noSelectRE.search(x), data))
        return labs

    def getBackedupLabels(self):
        """Return list of labels that should be backed up."""
        labs = self.getLabels()
        return filter(lambda x: not noBkupMailboxRE.search(x), labs)

    def getNonGmailLabels(self):
        """Return list of non-Gmail-permanent labels."""
        labs = self.getLabels()
        return filter(lambda x: not gmailMailboxRE.search(x), labs)

    def isLabelValid(self, label):
        """Return True iff 'label' is valid on server."""
        typ,data = self.s.status(label, '(UIDNEXT)')
        # Note that we don't call self._check('STATUS', typ, data)
        return typ == 'OK'

    def addLabel(self, uid, label):
        """Add 'label' to message with 'uid'."""
        self._selectAll()
        typ,data = self.s.uid('COPY', str(uid), label)
        self._check('UID COPY', typ, data)

    def removeLabel(self, uid, label):
        """Remove 'label' from message with 'uid'."""
        raise Exception('ERROR: removeLabel() is not implemented')

    def createLabel(self, label):
        """Create a new label 'label'."""
        if not self.isLabelValid(label):
            typ,data = self.s.create(label)
            self._check('CREATE', typ, data)

    def destroyLabel(self, label):
        """Remove label 'label' from all messages."""
        typ,data = self.s.delete(label)
        self._check('DELETE', typ, data)

    def deleteMessages(self, uids):
        """Delete message with 'uid'.  Leave in trash folder."""
        self._selectAll()
        if len(uids) > 0:
            arg = ",".join(map(lambda x: str(x), uids))
            typ,data = self.s.uid('COPY', arg, self.trashMail)
            self._check('UID COPY', typ, data)
            typ,data = self.s.uid('STORE', arg, '+FLAGS', r'(\Deleted)')
            self._check('UID STORE', typ, data)
            typ,data = self.s.expunge()
            self._check('EXPUNGE', typ, data)

    def getUIDs(self):
        """Return list of uids for all messages."""
        self._selectAll()
        typ,data = self.s.uid('SEARCH', None, 'ALL')
        self._check('UID SEARCH', typ, data)
        return map(lambda x: int(x), data[0].split())

    def getRecentUIDs(self, uid):
        """Return list of uids for all messages at or beyond 'uid'."""
        self._selectAll()

        # Convert UID to message sequence number.
        typ,data = self.s.uid('FETCH', '%d:*' % uid, '(INTERNALDATE)')
        self._check('FETCH', typ, data)
        if len(data) == 0 or data[0] == None:
            return []
        seqNo = int(data[0].split()[0])

        # Get UIDs of messages at or beyond this message sequence number.
        typ,data = self.s.uid('SEARCH', None, '%d:*' % seqNo)
        self._check('UID SEARCH', typ, data)
        return map(lambda x: int(x), data[0].split())

    def getUIDsWithLabel(self, label):
        """Return list of uids associated with 'label'."""
        typ,data = self.s.select(label)
        self._check('SELECT', typ, data)
        typ,data = self.s.uid('SEARCH', None, 'ALL')
        self._check('UID SEARCH', typ, data)
        return map(lambda x: int(x), data[0].split())

    def getSignatures(self, uids, selectAll=True):
        """Return list of (uid,signature) tuples.  A signature is
        a (date, identifier) pair.  The date is the date/time at which the
        message was received at Google; the identifier is the
        'Message-Id' message header if it exists (and is not too long)
        or the md5 hash of all the headers, otherwise.
        """
        if selectAll:
            self._selectAll()
        if len(uids) > 0:
            arg = ",".join(map(lambda x: str(x), uids))
            typ,data = self.s.uid('FETCH', arg,
                                  '(INTERNALDATE BODY.PEEK[HEADER])')
            self._check('UID FETCH', typ, data)
            data = filter(lambda x: isinstance(x, tuple), data)
            sigs = map(self._genSig, data)
            return sigs
        else:
            return []

    def getMsgs(self, uids):
        """Return a list of (uid,msg) pairs for messages 'uids'."""
        self._selectAll()
        if len(uids) > 0:
            arg = ",".join(map(lambda x: str(x), uids))
            typ,data = self.s.uid('FETCH', arg, '(UID BODY.PEEK[])')
            self._check('UID FETCH', typ, data)
            data = filter(lambda x: isinstance(x, tuple), data)
            return map(self._genMsg, data)
        else:
            return []

    def addMsg(self, datestr, msg):
        """Add a new message 'msg' with internal date 'datestr'."""
        typ,data = self.s.append(self.allMail, '\\Seen', datestr, msg)
        self._check('APPEND', typ, data)
        uid = int(appendRE.match(data[0]).group(1))
        return uid
    
    def disconnect(self):
        """Disconnect from IMAP server."""
        self.s.close()
        self.s.logout()

    #
    # Private methods
    #

    def _selectAll(self):
        """Select the '[Gmail]/All Mail' mailbox."""
        if self.currentMailbox != self.allMail:
            typ,data = self.s.select(self.allMail)
            self._check('SELECT', typ,data)
            self.currentMailbox = self.allMail

    def _genSig(self, item):
        """Generate a (uid,signature) pair from 'UID FETCH' response."""
        part1 = item[0]
        headers = item[1]
        headerLines = headers.splitlines()
        idate = idateRE.match(part1).group(1)
        date = time.strftime(self.dtFormat, imaplib.Internaldate2tuple(idate))
        uid = int(uidRE.match(part1).group(1))
        msgIdLines = filter(lambda x: msgIdRE.search(x), headerLines)
        id = ''
        if len(msgIdLines) > 0:
            id = msgIdRE.match(msgIdLines[0]).group(1)
        # Use hash instead of Message ID if (i) there is no ID or (ii) the
        # ID is too long.  The latter is necessary because the ID will be
        # used in a filename and some file systems (e.g., ext4) limit
        # file names to 255 or fewer characters.
        if len(id) > 0 and len(id) < 223:
            msgId = 'msgId-' + id
        else:
            m = hashlib.md5(headers)
            msgId = 'md5-' + m.hexdigest()
        return (uid, (date, msgId))

    def _genMsg(self, item):
        """Return a (uid,msg) pair from 'UID FETCH' response."""
        part1 = item[0]
        msg = item[1]
        uid = int(uidRE.match(part1).group(1))
        return (uid, msg)

    def _check(self,command, typ, data):
        """Confirm that 'typ' is 'OK'"""
        if typ != 'OK':
            raise Exception("ERROR processing '%s' command: '%s'" % (command,
                                                                     data[0]))


class gmailBackup:
    """Class for manager the backup of Gmail accounts.

    Instantiate with: gmailBackup(user, password, dir, logger, [retrieveLimit])

      user - Gmail account user name (e.g., user@gmail.com)
      password - Gmail account password
      dir - backup directory
      logger - logger object (defined in this module)
      retrieveLimit - max msgs pulled from IMAP server at once (default: 500)
    """

    dummySig = 'dummy-signature'
    indexFilesPattern = 'index_??????????????_????????_????????'

    #
    # Public methods
    #

    def __init__(self, user, password, dir, logger, retrieveLimit=500):
        self.sigs = None
        self.uids = None
        self.l = logger
        self.dir = dir
        self.retrieveLimit = retrieveLimit
        self.gmc = gmailConnection(user, password, logger)
        cn,cv,cd = self._getBackupUidNV(dir)
        self._setCurUIDNV(cn, cv, cd)
        self._cleanBackup()

    def getStats(self):
        """Return bytes sent and received since instantiation."""
        return self.gmc.getStats()

    def backup(self, disablePruning=False, disableLabelUpdate=True):
        """Perform backup."""
        self._checkVersion(init=True)
        nn,nv = self.gmc.getUidNV()
        self._setNewUIDNV(nn, nv)

        if self.curUIDValidity != self.newUIDValidity:
            self.l.l(self.l.VLEV_EXTRAS, 'UIDValidity has changed')

        sigs,mc = self._getNewMsgs()
        ma = len(sigs)

        if not disableLabelUpdate:
            bkupIndex,lb = self._getLabels()
        else:
            self.l.l(self.l.VLEV_STEPS, 'Not updating label index')
            bkupIndex = self._readIndexFile()
            lb = len(self._getBkupLabels(bkupIndex))
            for sig in sigs:
                if bkupIndex.get(sig) == None:
                    bkupIndex[sig] = []
        mb = len(bkupIndex) - 1
        self._writeIndexFile(self.newIndexFile + '.tmp', bkupIndex)
                
        md = 0
        if not disablePruning and not disableLabelUpdate:
            md = self._pruneBkupMsgs(bkupIndex)
        else:
            self.l.l(self.l.VLEV_STEPS, 'Not pruning backed up messages deleted on server')

        if os.path.isfile(self.curIndexFile):
            os.rename(self.curIndexFile, self.curIndexFile + '.bak')
        os.rename(self.newIndexFile + '.tmp', self.newIndexFile)
        if os.path.isfile(self.curIndexFile + '.bak'):
            os.remove(self.curIndexFile + '.bak')
        self._setCurUIDNV(self.newUIDNext, self.newUIDValidity,
                          self.newDate)
        return ma,md,mc,mb,lb


#
# I decided to comment this out, because it shouldn't really be necessary
# to blow away an account completely.  Uncomment it (and the call site,
# below) if you really want it.  Don't say I didn't warn you.
#

#     def zap(self):
#         """Delete all messages and labels in account on server."""
#         self.l.l(self.l.VLEV_STEPS, 'Zapping all mail in account...')
#         self.l.l(self.l.VLEV_STEPS, '  Deleting labels on server...')
#         labels = self.gmc.getNonGmailLabels()
#         self.l.l(self.l.VLEV_STEPS, '    Deleting %d labels' % len(labels))
#         for label in labels:
#             self.l.l(self.l.VLEV_PER_MSG, "    '%s'" % label)
#             self.gmc.destroyLabel(label)
#         uids = self._getUIDs()
#         self.l.l(self.l.VLEV_STEPS, '  Deleting messages on server...')
#         self.l.l(self.l.VLEV_STEPS, '    Deleting %d messages' % len(uids))
#         md = self.gmc.deleteMessages(uids)
#         return len(uids),len(labels)

    def restore(self):
        """Restore messages and labels from backup to server."""
        self._checkVersion()
        self.l.l(self.l.VLEV_STEPS, 'Restoring from backup...')
        bkupIndex = self._readIndexFile()
        self._augmentIndex(bkupIndex)
        la = self._restoreLabels(bkupIndex)
        self._restoreMsgs(bkupIndex)
        ma = len(bkupIndex) - 1
        return ma,la

    def disconnect(self):
        """Disconnect from server."""
        self.gmc.disconnect()

    #
    # Private methods
    #

    def _getBackupUidNV(self, dir):
        """Return UIDNEXT, UIDVALIDITY, and date associated with backup."""
        files = os.listdir(dir)
        indexFiles = fnmatch.filter(files, self.indexFilesPattern)
        bakIndexFiles = fnmatch.filter(files, self.indexFilesPattern + '.bak')
        date,uidNext,uidValidity,ext = None, 0, 0, ''
        file = None
        if len(indexFiles) > 0:
            indexFiles.sort(reverse=True)
            file = indexFiles[0]
        elif len(bakIndexFiles) > 0:
            bakIndexFiles.sort(reverse=True)
            file = bakIndexFiles[0]
        if file != None:
            dateStr,uidNext,uidValidity,ext = indexFileRE.match(file).groups()
            uidNext = int(uidNext, 16)
            uidValidity = int(uidValidity, 16)
            date = time.strptime(dateStr, self.gmc.dtFormat)
            if not os.path.isdir(self._getMailPath()):
                raise Exception('Backup directory corrupted')
            if ext == '.bak':
                self.l.l(self.l.VLEV_STEPS,
                         "Using old backup file ('%s')" % (file))
                os.rename(self._getPath(file),
                          self._getIndexFile(uidNext, uidValidity, date))
        return uidNext,uidValidity,date

    def _cleanBackup(self):
        """Prune old files (if any) in backup."""
        files = os.listdir(self.dir)
        indexFiles = filter(lambda x: (indexFileRE.search(x) and
                                       x != os.path.basename(self.curIndexFile)),
                           files)
        for file in indexFiles:
            os.remove(self._getPath(file))
            self.l.l(self.l.VLEV_STEPS, "Removing old file '%s'" % (file))
        # Perhaps we should also remove mail_ dirs
        # Perhaps we should remove .tmp files in mail_ dir

    def _getNewMsgs(self):
        """Retrieve new messages from server; store in backup."""
        self.l.l(self.l.VLEV_STEPS, 'Retrieving new messages from server...')
        if self.curUIDNext == self.newUIDNext:
            self.l.l(self.l.VLEV_STEPS, '  No new messages on server')
            return [],0
        if self.newUIDValidity == self.curUIDValidity:
            uids = self.gmc.getRecentUIDs(self.curUIDNext)
            sigs = self._getSignatures(uids)
        else:
            sigs = self._getSigs()
        prunedSigs = filter(self._getMsgLocally, sigs)
        sigsDict = dict(prunedSigs)
        
        self.l.l(self.l.VLEV_STEPS,
                 '  Retrieving %d new messages' % len(prunedSigs))
        prunedUids = [uid for uid,(date,msgId) in prunedSigs]
        for i in range(0, len(prunedUids), self.retrieveLimit):
            self._getMsgs(prunedUids[i:i+self.retrieveLimit], sigsDict)
        return prunedSigs,len(sigs)
    
    def _getMsgs(self, uids, sigsDict):
        """Retrieve specified messages from server."""
        results = self.gmc.getMsgs(uids)
        for uid,msg in results:
            date,msgId = sigsDict[uid]
            emlFile = self._getEmlFilePath(date, msgId)
            self.l.l(self.l.VLEV_PER_MSG,
                     "  Storing '%s'" % os.path.basename(emlFile))
            file = open(emlFile + '.tmp', 'w')
            file.write(msg)
            file.close()
            os.rename(emlFile + '.tmp', emlFile)

    def _getMsgLocally(self, sig):
        """Attempt to find msg with appropriate signature locally.  Return
        False if msg needs to be retrieved from server."""
        uid,(dateStr,msgId) = sig
        emlFile = self._getEmlFilePath(dateStr, msgId)
        if os.path.isfile(emlFile):
            return False
        else:
            return True

    def _getLabels(self):
        """Return backupIndex dictionary and number of labels on server."""
        self.l.l(self.l.VLEV_STEPS, 'Retrieving labels for each message...')
        labels = self.gmc.getBackedupLabels()
        self.l.l(self.l.VLEV_STEPS, '  Retrieving %d labels' % len(labels))
        allSigs = self._getSigs()
        bkupIndex = dict([(sig,[]) for uid,sig in allSigs])
        bkupIndex[(None,None)] = []
        for l in labels:
            uids = self.gmc.getUIDsWithLabel(l)
            self.l.l(self.l.VLEV_PER_MSG,
                     "  '%s' (%d messages)" % (l, len(uids)))
            if len(uids) > 0:
                sigs = self._getSignatures(uids, selectAll=False)
            else:
                sigs = [(None,(None,None))]
            for sig in sigs:
                u,(dateStr,msgId) = sig
                if bkupIndex.get((dateStr,msgId)) != None:
                    bkupIndex[(dateStr,msgId)].append(l)
        return bkupIndex,len(labels)

    def _pruneBkupMsgs(self, bkupIndex):
        """Prune messages in backup if they are not in 'bkupIndex'. Return
        number of messages deleted."""
        self.l.l(self.l.VLEV_STEPS, 'Removing messages no longer on server...')
        files = os.listdir(self._getMailPath())
        msgCount = 0
        for file in files:
            base,dateStr,msgId = emlFileRE.match(file).groups()
            if bkupIndex.get((dateStr,msgId)) == None:
                self.l.l(self.l.VLEV_PER_MSG, "  Removing '%s'" % (file))
                os.remove(self._getMailPath() + os.sep + file)
                msgCount += 1
        self.l.l(self.l.VLEV_STEPS, '  %d messages removed' % msgCount)
        return msgCount

    def _getBkupLabels(self, bkupIndex):
        labelSet = set()
        for sig,labels in bkupIndex.iteritems():
            for label in labels:
                labelSet.add(label)
        return labelSet
    
    def _restoreLabels(self, bkupIndex):
        """Restore labels appearing in 'bkupIndex' to the server."""
        self.l.l(self.l.VLEV_STEPS, '  Restoring labels...')
        labelSet = self._getBkupLabels(bkupIndex)
        restorableLabels = filter(lambda x: not gmailMailboxRE.search(x),
                                  labelSet)
        self.l.l(self.l.VLEV_STEPS,
                 '    Restoring %d labels' % len(restorableLabels))
        for label in restorableLabels:
            self.l.l(self.l.VLEV_PER_MSG, "    '%s'" % label)
            self.gmc.createLabel(label)
        return len(restorableLabels)

    def _restoreMsgs(self, bkupIndex):
        """Restore to server all messages appearing in 'bkupIndex'."""
        self.l.l(self.l.VLEV_STEPS, '  Restoring messages...')
        self.l.l(self.l.VLEV_STEPS,
                 '    Restoring %d messages' % (len(bkupIndex) - 1))
        for sig,labels in bkupIndex.iteritems():
            if sig != (None,None):
                date,msgId = sig
                self.l.l(self.l.VLEV_PER_MSG,
                         '    %s' % self._sigStr(date, msgId), newLine=False)
                newUid = self._restoreMsg(sig, labels)

    def _restoreMsg(self, sig, labs):
        dateStr,msgId = sig
        emlFile = self._getEmlFilePath(dateStr, msgId)
        file = open(emlFile, 'r')
        msg = file.read()
        file.close()
        date = time.strptime(dateStr, self.gmc.dtFormat)
        iDateStr = imaplib.Time2Internaldate(time.mktime(date))
        uid = self.gmc.addMsg(iDateStr, msg)
        self.l.l(self.l.VLEV_PER_MSG,
                 '[%s]' % ' '.join(map(lambda x: "'" + x + "'", labs)))
        for l in labs:
            self.gmc.addLabel(uid, l)
        return uid
        
    def _readIndexFile(self):
        if os.path.isfile(self.curIndexFile):
            file = open(self.curIndexFile, 'r')
            lines = file.read().splitlines()
            return dict(map(lambda x: eval(x), lines))
        else:
            return {}

    def _writeIndexFile(self, file, bkupIndex):
        f = open(file, 'w')
        for sig,labels in bkupIndex.iteritems():
            item = (sig, labels)
            f.write(str(item) + '\n')
        f.close()

    def _augmentIndex(self, bkupIndex):
        """Ensure that all the messages in the backup are in the index."""
        if bkupIndex.get((None,None)) == None:
            bkupIndex[(None,None)] = []
        files = os.listdir(self._getMailPath())
        for file in files:
            base,dateStr,msgId = emlFileRE.match(file).groups()
            if bkupIndex.get((dateStr,msgId)) == None:
                bkupIndex[(dateStr,msgId)] = []

    def _checkVersion(self, init=False):
        if os.path.isfile(self._getVersionFile()):
            formatVer = version.split('-')[1].split('.')
            bkupFormatVer = self._getBackupVersion().split('-')[1].split('.')
            if (int(formatVer[0]) != int(bkupFormatVer[0]) or
                int(formatVer[1]) != int(bkupFormatVer[1])):
                raise Exception('Backup storage format version (' + self._getBackupVersion() +
                                ') is incompatible with this version (' +
                                version + ') of yagbut.')
            self._setBackupVersion()
        elif init:
            self._setBackupVersion()
        else:
            raise Exception('Backup appears to be corrupted (no version file)')

    def _getBackupVersion(self):
        f = open(self._getVersionFile(), 'r')
        ver = f.read().splitlines()[0]
        f.close()
        return ver

    def _setBackupVersion(self):
        tmpVersionFile = self._getVersionFile() + '.tmp'
        versionFile = self._getVersionFile()
        f = open(tmpVersionFile, 'w')
        f.write(version + '\n')
        f.close()
        if os.path.isfile(versionFile):
            os.remove(versionFile)
        os.rename(tmpVersionFile, versionFile)

    def _getVersionFile(self):
        return self._getPath('version')

    def _getSignatures(self, uids, selectAll=True):
        sigs = []
        for i in range(0, len(uids), self.retrieveLimit):
            sigs += self.gmc.getSignatures(uids[i:i+self.retrieveLimit],
                                           selectAll)
        return map(lambda (id,(d,i)): (id,(d,string.replace(i, os.sep, ''))), sigs)

    def _getSigs(self):
        if self.sigs == None:
            self.sigs = self._getSignatures(self._getUIDs())
        return self.sigs

    def _getUIDs(self):
        if self.uids == None:
            self.uids = self.gmc.getUIDs()
        return self.uids

    def _setCurUIDNV(self, uidNext, uidValidity, date):
        self.curUIDNext = uidNext
        self.curUIDValidity = uidValidity
        self.curIndexFile = self._getIndexFile(uidNext, uidValidity, date)
        self.sigs = None
        self.uids = None

    def _setNewUIDNV(self, uidNext, uidValidity):
        self.newUIDNext = uidNext
        self.newUIDValidity = uidValidity
        self.newDate = time.localtime()
        self.newIndexFile = self._getIndexFile(uidNext, uidValidity,
                                             self.newDate)
        if not os.path.isdir(self._getMailPath()):
            os.mkdir(self._getMailPath())

    def _sigStr(self, dateStr, msgId):
        return dateStr + '_' + msgId

    def _getPath(self, name):
        return self.dir + os.sep + name

    def _getMailPath(self):
        return self._getPath('mail')

    def _getIndexFile(self, uidNext, uidValidity, date):
        if date == None:
            dateStr = '00000000000000'
        else:
            dateStr = time.strftime(self.gmc.dtFormat, date)
        return self._getPath('index_%s_%08x_%08x' % (dateStr, uidNext,
                                                     uidValidity))

    def _getEmlFile(self, dateStr, msgId):
        return self._sigStr(dateStr, msgId) + '.eml'

    def _getEmlFilePath(self, dateStr, msgId):
        prefix = self._getMailPath() + os.sep
        return prefix + self._getEmlFile(dateStr, msgId)

class logger:
    VLEV_ERRS = 0     # Nothing but erros
    VLEV_SUMMARY = 1  # Just summary info
    VLEV_STEPS = 2    # Info about steps being performed
    VLEV_PER_MSG = 3  # Per msg/label info
    VLEV_EXTRAS = 4   # Everything
    level = VLEV_STEPS

    def set(self, lev):
        self.level = lev

    def l(self, lev, str, newLine=True):
        if lev <= self.level:
            print str,
            if newLine:
                print

###
### REMAINING FUNCTIONS
###

#
# usage -
#
#   Print usage information.
#

def usage():
    print 'yagbut version %s' % version
    print '''Usage: yagbut.py [OPTION]...
Backup, restore, or zap a Gmail mail account.

At least one of the following options must be specified.
  -b, --backup               Backup Gmail account to local directory.  Cannot
                             be used --restore nor --zap.
  -r, --restore              Restore Gmail account from local directory.
                             Cannot be used with --backup.
  -z, --zap                  Erase messages on Gmail account.  Cannot be used
                             with -backup.  USE WITH CAUTION!!!

Other options.
  -d, --dir DIR               Backup to directory/folder DIR (default: .).
  -l, --disable-label-update  Don't update label index (implies -D);
                              can only be used with --backup.
  -D, --disable-prune         Don't remove deleted messages in backup;
                              can only be used with --backup.
  -F, --password-file FILE    Use password in FILE; cannot use with --password.
  -h, --help                  Print usage.
  -H, --host HOST             Connect to HOST (default: imap.googlemail.com)
  -p, --password PASSWORD     Use password PASSWORD; cannot use with
                              --password-file.
  -r, --retrieve-limit LIMIT  Only grab LIMIT messages from server at a time
                              (default: 500); can only be used with --backup.
  -u, --user USER             Use USER (e.g., myuser@gmail.com)
  -v, --verbose LEVEL         Set verbosity level (default: 2)
                                0 Only errors
                                1 Summary info
                                2 Step-by-step info
                                3 Per-message info
                                4 Everything

Exit status
0 if OK,
1 if error.

Examples.

To backup:

  % yagbut.py --backup --dir /home/myaccount/bkups/gmail-myaccount --user myaccount@gmail.com

To restore:

  % yagbut.py --zap --restore --dir /home/myaccount/bkups/gmail-myaccount --user myaccount@gmail.com

Warning: --zap destroys all mail on the Gmail server in the specified account!

Report bugs at http://code.google.com/p/yagbut/issues/list.''',

#
# main -
#
#   Main entry point.
#

def main():
    defaultHostName = 'imap.googlemail.com'

    locale.setlocale(locale.LC_ALL, "")
    l = logger()

    try:
        opts,args = getopt.getopt(sys.argv[1:], 'hbzrH:u:p:v:Dd:F:r:l',
                                  ['help', 'backup', 'zap', 'restore',
                                   'host=', 'user=', 'password=', 'verbose=',
                                   'disable-prune', 'dir=', 'password-file=',
                                   'retrieve-limit=', 'disable-label-update'])
    except getopt.GetoptError, err:
        print str(err)
        print "Use --help for usage information."
        sys.exit(2)
    performingBackup = False
    performingZap = False
    performingRestore = False
    hostName = defaultHostName
    userName = None
    password = None
    disPrune = False
    disLabUp = False
    passwordFile = None
    retrieveLimit = None
    logLevel = 0
    dir = '.'
    for o,a in opts:
        if o in ('-v', '--verbose'):
            logLevel = int(a)
            l.set(logLevel)
        elif o in ('-h', '--help'):
            usage()
            sys.exit(0)
        elif o in ('-b', '--backup'):
            performingBackup = True
        elif o in ('-z', '--zap'):
            performingZap = True
        elif o in ('-r', '--restore'):
            performingRestore = True
        elif o in ('-H', '--host'):
            hostName = a
        elif o in ('-u', '--user'):
            userName = a
        elif o in ('-p', '--password'):
            password = a
        elif o in ('-D', '--disable-prune'):
            disPrune = True
        elif o in ('-l', '--disable-label-update'):
            disLabUp = True
        elif o in ('-d', '--dir'):
            dir = os.path.dirname(a + os.sep)
        elif o in ('-F', '--password-file'):
            passwordFile = os.path.abspath(a)
        elif o in ('-r', '--retrieve-limit'):
            retrieveLimit = int(a)

    dir = os.path.abspath(dir)
    earlyTermination = False
    if not performingBackup and not performingZap and not performingRestore:
        print 'At least one of --backup, --zap, or --restore must be specified.'
        earlyTerimantion = True
    if performingBackup and (performingZap or performingRestore):
        print "Can't use --backup with --zap or --restore."
        earlyTermination = True
    if userName == None:
        print 'Must specify user name (--user).'
        earlyTermination = True
    if password != None and passwordFile != None:
        print 'Password cannot be specified via both --password and --password-file.'
        earlyTermination = True
    if disPrune and not performingBackup:
        print '--disable-prune is meaningless when backup is not being performed.'
        earlyTermination = True
    if disLabUp and not performingBackup:
        print '--disable-label-update is meaningless when backup is not being performed.'
        earlytermination = True
    if retrieveLimit != None and not performingBackup:
        print '--retrieve-limit is meaningless when backup is not being performed.'
        earlyTermination = True
    if retrieveLimit != None and retrieveLimit < 1:
        print 'Only positive values should be used with --retrieve-limit.'
        earlyTermination = True
    if not os.path.isdir(dir):
        print "Backup directory '%s' does not exist." % (dir)
        earlyTermination = True
    if passwordFile != None and not os.path.isfile(passwordFile):
        print "Password file '%s' does not exist." % (passwordFile)
        earlyTermination = True
    if logLevel < l.VLEV_ERRS or logLevel > l.VLEV_EXTRAS:
        print "Verbosity level of '%d' is out of range." % (logLevel)
        earlyTermination = True

    if earlyTermination:
        print "Use --help for usage information."
        sys.exit(1)

    l.l(l.VLEV_STEPS, "Backup directory: '%s'" % dir)
        
    if performingZap:
        print 'Account: ' + userName
        print 'Host: ' + hostName
        print 'You about to zap everything for account above.'
        response = raw_input('Are you certain that this is what you want to do (yes/no)? ')
        if response.lower() != 'yes':
            sys.exit()

    if passwordFile != None:
        file = open(passwordFile, 'r')
        password = file.read().splitlines()[0]
        print "Password read frome file '%s'" % passwordFile
        file.close()
    if password == None:
        password = getpass.getpass('Password: ')

    if retrieveLimit == 0 or retrieveLimit == None:
        retrieveLimit = 500

    startTime = time.time()

    try:
        bu = gmailBackup(userName, password, dir, l, retrieveLimit)
    except Exception,e:
        print "ERROR:", e
        sys.exit(1)

    try:
        if performingBackup:
            ma,md,mc,mb,lb = bu.backup(disablePruning=disPrune,
                                       disableLabelUpdate=disLabUp)
        if performingZap:
            print "!!!!!! Zap functionality has been commented out.  If you"
            print "!!!!!! really want it, you'll have to edit the code."
            md,ld = 0,0
#            md,ld = bu.zap()
        if performingRestore:
            ma,la = bu.restore()
    
    finally:
        bu.disconnect()

    totalTime = time.time() - startTime

    l.l(l.VLEV_SUMMARY, 'Backup summary...')
    if performingBackup:
        l.l(l.VLEV_SUMMARY, '  Messages added to backup: %d' % ma)
        l.l(l.VLEV_SUMMARY, '  Messages deleted from backup: %d' % md)
        l.l(l.VLEV_SUMMARY, '  Messages considered: %d' % mc)
        l.l(l.VLEV_SUMMARY, '  Messages in backup: %d' % mb)
        l.l(l.VLEV_SUMMARY, '  Labels in backup: %d' % lb)
    if performingZap:
        l.l(l.VLEV_SUMMARY, '  Messages deleted from server: %d' % md)
        l.l(l.VLEV_SUMMARY, '  Labels deleted from server: %d' % ld)
    if performingRestore:
        l.l(l.VLEV_SUMMARY, '  Messages restored to server: %d' % ma)
        l.l(l.VLEV_SUMMARY, '  Labels restored to server: %d' % la)
    bytesSent, bytesRecv = bu.getStats()
    bsStr = locale.format('%d', bytesSent, True)
    brStr = locale.format('%d', bytesRecv, True)
    l.l(l.VLEV_SUMMARY, '  Bytes sent: %s' % bsStr)
    l.l(l.VLEV_SUMMARY, '  Bytes recv: %s' % brStr)
    l.l(l.VLEV_SUMMARY, '  Total processing time: %s' %
        datetime.timedelta(seconds=totalTime))

if __name__ == '__main__':
    main()
    sys.exit(0)
